package org.zjx.analysis.identify;

import lombok.SneakyThrows;
import net.sourceforge.tess4j.Tesseract;
import org.zjx.command.util.ImageUtil;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class IdentifyTessUtil {
    // 缓存配置，避免重复创建相同的Tesseract实例
    private static final ConcurrentHashMap<String, ThreadLocal<Tesseract>> TESSERACT_CACHE = new ConcurrentHashMap<>();


    /**
     * 批量识别多个区域
     */
    public static Map<String, String> batchIdentify(BufferedImage sourceImage, Map<String, Rectangle> areas) {
        Map<String, String> results = new ConcurrentHashMap<>();

        // 并行处理所有区域
        areas.entrySet().parallelStream().forEach(entry -> {
            String areaName = entry.getKey();
            Rectangle rect = entry.getValue();

            BufferedImage subImage = ImageUtil.cutImage(sourceImage, rect);
            String result = tessIdentify(subImage);

            results.put(areaName, result);
        });

        return results;
    }

    /**
     * 批量识别多个区域（指定语言和白名单）
     */
    public static Map<String, String> batchIdentify(BufferedImage sourceImage,
                                                    Map<String, Rectangle> areas,
                                                    String language,
                                                    String charWhitelist) {
        Map<String, String> results = new ConcurrentHashMap<>();

        areas.entrySet().parallelStream().forEach(entry -> {
            String areaName = entry.getKey();
            Rectangle rect = entry.getValue();

            BufferedImage subImage = ImageUtil.cutImage(sourceImage, rect);
            String result = tessIdentifyWithLanguage(subImage, language, charWhitelist);

            results.put(areaName, result);
        });

        return results;
    }



    /**
     * 主要的识别方法，支持中英文混合识别（使用默认白名单 "0123456789/"）
     */
    public static String tessIdentify(BufferedImage bufferedImage) {
        String resultChi = tessIdentifyWithLanguage(bufferedImage, "chi_sim", "/0123456789");
        String resultEng = tessIdentifyWithLanguage(bufferedImage, "eng", "/0123456789");

        return mergeAndCorrectResults(resultChi, resultEng);
    }

    /**
     * 单独调用指定语言的识别方法
     *
     * @param bufferedImage 待识别图片
     * @param language      语言 ("chi_sim" 或 "eng")
     * @param charWhitelist 字符白名单
     * @return 识别结果
     */
    public static String tessIdentify(BufferedImage bufferedImage, String language, String charWhitelist) {
        return tessIdentifyWithLanguage(bufferedImage, language, charWhitelist);
    }

    /**
     * 逐个字符比对并修正识别结果
     */
    private static String mergeAndCorrectResults(String chi, String eng) {
        if (chi == null) chi = "";
        if (eng == null) eng = "";

        StringBuilder finalResult = new StringBuilder();
        int i = 0, j = 0;

        while (i < chi.length() || j < eng.length()) {
            char chiChar = (i < chi.length()) ? chi.charAt(i) : '\0';
            char engChar = (j < eng.length()) ? eng.charAt(j) : '\0';

            // 情况1: 英文识别为'/'，中文识别为'1' → 采用'/'
            if (engChar == '/' && chiChar == '1') {
                finalResult.append('/');
                i++;
                j++;
            }
            // 情况2: 中文识别为'0'，英文不是'0'，且英文字符在中文'0'之后 → 添加'0'，中文指针前进
            else if (chiChar == '0' && engChar != '0' && engChar != '\0') {
                finalResult.append('0');
                i++;
            }
            // 情况3: 英文识别为数字或斜杠，优先采用英文结果
            else if (isValidChar(engChar)) {
                finalResult.append(engChar);
                j++;
                // 如果中文当前字符也是有效字符，但被英文跳过，中文指针也前进
                if (isValidChar(chiChar)) {
                    i++;
                }
            }
            // 情况4: 英文无效但中文有效，采用中文结果
            else if (isValidChar(chiChar)) {
                finalResult.append(chiChar);
                i++;
                // 如果英文当前字符也是有效字符，但被中文跳过，英文指针也前进
                if (isValidChar(engChar)) {
                    j++;
                }
            }
            // 情况5: 两个都无效，跳过当前字符
            else {
                if (chiChar != '\0') i++;
                if (engChar != '\0') j++;
            }
        }

        return finalResult.toString();
    }

    /**
     * 检查字符是否为有效字符（数字或斜杠）- 用于默认中英文比较
     */
    private static boolean isValidChar(char c) {
        return (c >= '0' && c <= '9') || c == '/' || c ==  ',';
    }

    /**
     * 使用指定语言进行 OCR 识别
     */
    @SneakyThrows
    private static String tessIdentifyWithLanguage(BufferedImage bufferedImage, String language, String charWhitelist) {
        Tesseract tesseract = getTesseract(language, charWhitelist);
        return tesseract.doOCR(bufferedImage);
    }

    /**
     * 获取指定语言的Tesseract实例（线程安全）
     */
    private static Tesseract getTesseract(String language, String charWhitelist) {
        String cacheKey = language + "|" + charWhitelist;
        return TESSERACT_CACHE
                .computeIfAbsent(cacheKey, key -> ThreadLocal.withInitial(() -> createTesseract(language, charWhitelist)))
                .get();
    }

    /**
     * 创建 Tesseract 实例
     */
    private static Tesseract createTesseract(String language, String charWhitelist) {
        Tesseract tesseract = new Tesseract();
        tesseract.setVariable("user_defined_dpi", "300");
        tesseract.setDatapath("src/main/resources/tessdata");
        if (charWhitelist != null)
            tesseract.setVariable("tessedit_char_whitelist", charWhitelist);
        tesseract.setLanguage(language);
        tesseract.setPageSegMode(7);
        tesseract.setOcrEngineMode(1);
        return tesseract;
    }
}
